package websocket

import (
	"net/http"
	"net/url"
	"sync"

	"github.com/gorilla/websocket"
)

type WebSocketProxy struct {
	upgrader    websocket.Upgrader
	backend     string
	readBuffer  int
	writeBuffer int
	mu          sync.RWMutex
}

func NewWebSocketProxy(backend string, readBuffer, writeBuffer int) *WebSocketProxy {
	return &WebSocketProxy{
		upgrader: websocket.Upgrader{
			ReadBufferSize:  readBuffer,
			WriteBufferSize: writeBuffer,
			CheckOrigin: func(r *http.Request) bool {
				return true // 可以根据需求设置更严格的检查
			},
		},
		backend:     backend,
		readBuffer:  readBuffer,
		writeBuffer: writeBuffer,
	}
}

func (wp *WebSocketProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// 升级HTTP连接到WebSocket
	clientConn, err := wp.upgrader.Upgrade(w, r, nil)
	if err != nil {
		http.Error(w, "WebSocket upgrade failed", http.StatusInternalServerError)
		return
	}
	defer clientConn.Close()

	// 解析后端URL
	backendURL, err := url.Parse(wp.backend)
	if err != nil {
		return
	}

	// 连接到后端WebSocket服务器
	backendConn, _, err := websocket.DefaultDialer.Dial(getWebSocketURL(*backendURL, r), nil)
	if err != nil {
		return
	}
	defer backendConn.Close()

	// 双向代理WebSocket消息
	errChan := make(chan error, 2)
	go wp.proxyWebSocket(clientConn, backendConn, errChan)
	go wp.proxyWebSocket(backendConn, clientConn, errChan)

	// 等待任一方向出错或连接关闭
	<-errChan
}

func (wp *WebSocketProxy) proxyWebSocket(dst, src *websocket.Conn, errChan chan error) {
	for {
		messageType, message, err := src.ReadMessage()
		if err != nil {
			errChan <- err
			return
		}

		err = dst.WriteMessage(messageType, message)
		if err != nil {
			errChan <- err
			return
		}
	}
}

func getWebSocketURL(backendURL url.URL, r *http.Request) string {
	scheme := "ws"
	if r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https" {
		scheme = "wss"
	}
	backendURL.Scheme = scheme
	return backendURL.String()
}

func (wp *WebSocketProxy) UpdateBackend(backend string) {
	wp.mu.Lock()
	defer wp.mu.Unlock()
	wp.backend = backend
}

// WebSocket连接追踪器
type ConnectionTracker struct {
	connections map[string]*websocket.Conn
	mu          sync.RWMutex
}

func NewConnectionTracker() *ConnectionTracker {
	return &ConnectionTracker{
		connections: make(map[string]*websocket.Conn),
	}
}

func (ct *ConnectionTracker) Add(id string, conn *websocket.Conn) {
	ct.mu.Lock()
	defer ct.mu.Unlock()
	ct.connections[id] = conn
}

func (ct *ConnectionTracker) Remove(id string) {
	ct.mu.Lock()
	defer ct.mu.Unlock()
	delete(ct.connections, id)
}

func (ct *ConnectionTracker) Get(id string) (*websocket.Conn, bool) {
	ct.mu.RLock()
	defer ct.mu.RUnlock()
	conn, ok := ct.connections[id]
	return conn, ok
}

func (ct *ConnectionTracker) CloseAll() {
	ct.mu.Lock()
	defer ct.mu.Unlock()
	for _, conn := range ct.connections {
		conn.Close()
	}
	ct.connections = make(map[string]*websocket.Conn)
}
